import os
import sys
from PyQt5.QtWidgets import (QWidget, QVBoxLayout, QHBoxLayout, QLabel, 
                             QPushButton, QFileDialog, QSpinBox, QComboBox,
                             QGroupBox, QGridLayout, QProgressBar, QMessageBox,
                             QSlider, QCheckBox, QLineEdit)
from PyQt5.QtCore import Qt, QThread, pyqtSignal, QTimer
from PyQt5.QtGui import QPixmap, QImage
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont


class LogoAdderThread(QThread):
    progress_updated = pyqtSignal(int)
    finished = pyqtSignal(str)
    error_occurred = pyqtSignal(str)

    def __init__(self, video_path, logo_path, logo_size, logo_position, output_path):
        super().__init__()
        self.video_path = video_path
        self.logo_path = logo_path
        self.logo_size = logo_size
        self.logo_position = logo_position
        self.output_path = output_path

    def run(self):
        try:
            # 读取视频
            cap = cv2.VideoCapture(self.video_path)
            if not cap.isOpened():
                self.error_occurred.emit("无法打开视频文件")
                return

            # 获取视频信息
            total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            fps = cap.get(cv2.CAP_PROP_FPS)
            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

            # 读取logo图片
            logo = cv2.imread(self.logo_path, cv2.IMREAD_UNCHANGED)
            if logo is None:
                self.error_occurred.emit("无法读取logo图片")
                return

            # 调整logo大小
            logo_height, logo_width = logo.shape[:2]
            scale = self.logo_size / max(logo_width, logo_height)
            new_width = int(logo_width * scale)
            new_height = int(logo_height * scale)
            logo = cv2.resize(logo, (new_width, new_height))

            # 创建视频写入器
            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
            out = cv2.VideoWriter(self.output_path, fourcc, fps, (width, height))

            # 计算logo位置
            if self.logo_position == "左上角":
                x, y = 10, 10
            elif self.logo_position == "右上角":
                x, y = width - new_width - 10, 10
            elif self.logo_position == "左下角":
                x, y = 10, height - new_height - 10
            elif self.logo_position == "右下角":
                x, y = width - new_width - 10, height - new_height - 10
            elif self.logo_position == "中心":
                x = (width - new_width) // 2
                y = (height - new_height) // 2
            else:
                x, y = 10, 10

            frame_count = 0
            while True:
                ret, frame = cap.read()
                if not ret:
                    break

                # 添加logo到帧
                if logo.shape[2] == 4:  # 有透明通道
                    # 处理透明logo
                    alpha = logo[:, :, 3] / 255.0
                    alpha = np.expand_dims(alpha, axis=2)
                    
                    # 确保坐标在有效范围内
                    y1, y2 = max(0, y), min(height, y + new_height)
                    x1, x2 = max(0, x), min(width, x + new_width)
                    
                    if y1 < y2 and x1 < x2:
                        roi = frame[y1:y2, x1:x2]
                        logo_roi = logo[y1-y:y2-y, x1-x:x2-x]
                        
                        if logo_roi.shape[:2] == roi.shape[:2]:
                            alpha_roi = alpha[y1-y:y2-y, x1-x:x2-x]
                            for c in range(3):
                                roi[:, :, c] = (1 - alpha_roi[:, :, 0]) * roi[:, :, c] + alpha_roi[:, :, 0] * logo_roi[:, :, c]
                else:
                    # 处理不透明logo
                    y1, y2 = max(0, y), min(height, y + new_height)
                    x1, x2 = max(0, x), min(width, x + new_width)
                    
                    if y1 < y2 and x1 < x2:
                        roi = frame[y1:y2, x1:x2]
                        logo_roi = logo[y1-y:y2-y, x1-x:x2-x]
                        
                        if logo_roi.shape[:2] == roi.shape[:2]:
                            roi[:] = logo_roi

                out.write(frame)
                frame_count += 1
                
                # 更新进度
                progress = int((frame_count / total_frames) * 100)
                self.progress_updated.emit(progress)

            cap.release()
            out.release()
            self.finished.emit(self.output_path)

        except Exception as e:
            self.error_occurred.emit(f"处理过程中出现错误: {str(e)}")


class LogoAdderWidget(QWidget):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.video_path = ""
        self.logo_path = ""
        self.output_path = ""
        self.thread = None
        self.init_ui()

    def init_ui(self):
        layout = QVBoxLayout()
        
        # 视频选择组
        video_group = QGroupBox("选择视频文件")
        video_layout = QHBoxLayout()
        self.video_label = QLabel("未选择视频文件")
        self.video_btn = QPushButton("选择视频")
        self.video_btn.clicked.connect(self.select_video)
        video_layout.addWidget(self.video_label)
        video_layout.addWidget(self.video_btn)
        video_group.setLayout(video_layout)
        
        # Logo选择组
        logo_group = QGroupBox("选择Logo图片")
        logo_layout = QHBoxLayout()
        self.logo_label = QLabel("未选择logo文件")
        self.logo_btn = QPushButton("选择Logo")
        self.logo_btn.clicked.connect(self.select_logo)
        logo_layout.addWidget(self.logo_label)
        logo_layout.addWidget(self.logo_btn)
        logo_group.setLayout(logo_layout)
        
        # Logo预览
        self.logo_preview = QLabel()
        self.logo_preview.setFixedSize(200, 150)
        self.logo_preview.setStyleSheet("border: 1px solid gray;")
        self.logo_preview.setAlignment(Qt.AlignCenter)
        self.logo_preview.setText("Logo预览")
        
        # Logo设置组
        settings_group = QGroupBox("Logo设置")
        settings_layout = QGridLayout()
        
        # Logo大小设置
        settings_layout.addWidget(QLabel("Logo大小:"), 0, 0)
        self.size_spinbox = QSpinBox()
        self.size_spinbox.setRange(10, 500)
        self.size_spinbox.setValue(100)
        self.size_spinbox.setSuffix(" px")
        settings_layout.addWidget(self.size_spinbox, 0, 1)
        
        # Logo位置设置
        settings_layout.addWidget(QLabel("Logo位置:"), 1, 0)
        self.position_combo = QComboBox()
        self.position_combo.addItems(["左上角", "右上角", "左下角", "右下角", "中心"])
        settings_layout.addWidget(self.position_combo, 1, 1)
        
        # 自定义位置
        self.custom_position_check = QCheckBox("自定义位置")
        self.custom_position_check.toggled.connect(self.toggle_custom_position)
        settings_layout.addWidget(self.custom_position_check, 2, 0)
        
        # X坐标
        settings_layout.addWidget(QLabel("X坐标:"), 3, 0)
        self.x_spinbox = QSpinBox()
        self.x_spinbox.setRange(0, 9999)
        self.x_spinbox.setValue(10)
        self.x_spinbox.setSuffix(" px")
        self.x_spinbox.setEnabled(False)
        settings_layout.addWidget(self.x_spinbox, 3, 1)
        
        # Y坐标
        settings_layout.addWidget(QLabel("Y坐标:"), 4, 0)
        self.y_spinbox = QSpinBox()
        self.y_spinbox.setRange(0, 9999)
        self.y_spinbox.setValue(10)
        self.y_spinbox.setSuffix(" px")
        self.y_spinbox.setEnabled(False)
        settings_layout.addWidget(self.y_spinbox, 4, 1)
        
        settings_group.setLayout(settings_layout)
        
        # 输出设置组
        output_group = QGroupBox("输出设置")
        output_layout = QHBoxLayout()
        self.output_label = QLabel("未设置输出路径")
        self.output_btn = QPushButton("选择输出路径")
        self.output_btn.clicked.connect(self.select_output)
        output_layout.addWidget(self.output_label)
        output_layout.addWidget(self.output_btn)
        output_group.setLayout(output_layout)
        
        # 进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setVisible(False)
        
        # 开始处理按钮
        self.process_btn = QPushButton("开始添加Logo")
        self.process_btn.clicked.connect(self.start_processing)
        self.process_btn.setEnabled(False)
        
        # 添加所有组件到主布局
        layout.addWidget(video_group)
        layout.addWidget(logo_group)
        layout.addWidget(self.logo_preview, alignment=Qt.AlignCenter)
        layout.addWidget(settings_group)
        layout.addWidget(output_group)
        layout.addWidget(self.progress_bar)
        layout.addWidget(self.process_btn)
        layout.addStretch()
        
        self.setLayout(layout)

    def select_video(self):
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择视频文件", "", 
            "视频文件 (*.mp4 *.avi *.mov *.mkv *.wmv *.flv *.webm)"
        )
        if file_path:
            self.video_path = file_path
            self.video_label.setText(os.path.basename(file_path))
            self.check_ready()

    def select_logo(self):
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择Logo图片", "", 
            "图片文件 (*.png *.jpg *.jpeg *.bmp *.gif *.tiff)"
        )
        if file_path:
            self.logo_path = file_path
            self.logo_label.setText(os.path.basename(file_path))
            self.load_logo_preview()
            self.check_ready()

    def load_logo_preview(self):
        if self.logo_path:
            pixmap = QPixmap(self.logo_path)
            if not pixmap.isNull():
                pixmap = pixmap.scaled(180, 130, Qt.KeepAspectRatio, Qt.SmoothTransformation)
                self.logo_preview.setPixmap(pixmap)
            else:
                self.logo_preview.setText("无法加载图片")

    def select_output(self):
        file_path, _ = QFileDialog.getSaveFileName(
            self, "选择输出文件", "", 
            "MP4文件 (*.mp4)"
        )
        if file_path:
            if not file_path.endswith('.mp4'):
                file_path += '.mp4'
            self.output_path = file_path
            self.output_label.setText(os.path.basename(file_path))
            self.check_ready()

    def toggle_custom_position(self, checked):
        self.position_combo.setEnabled(not checked)
        self.x_spinbox.setEnabled(checked)
        self.y_spinbox.setEnabled(checked)

    def check_ready(self):
        ready = bool(self.video_path and self.logo_path and self.output_path)
        self.process_btn.setEnabled(ready)

    def start_processing(self):
        if not all([self.video_path, self.logo_path, self.output_path]):
            QMessageBox.warning(self, "警告", "请先选择视频、logo和输出路径")
            return

        # 获取设置
        logo_size = self.size_spinbox.value()
        
        if self.custom_position_check.isChecked():
            logo_position = f"自定义({self.x_spinbox.value()},{self.y_spinbox.value()})"
        else:
            logo_position = self.position_combo.currentText()

        # 禁用按钮
        self.process_btn.setEnabled(False)
        self.progress_bar.setVisible(True)
        self.progress_bar.setValue(0)

        # 创建并启动处理线程
        self.thread = LogoAdderThread(
            self.video_path, self.logo_path, logo_size, logo_position, self.output_path
        )
        self.thread.progress_updated.connect(self.update_progress)
        self.thread.finished.connect(self.processing_finished)
        self.thread.error_occurred.connect(self.processing_error)
        self.thread.start()

    def update_progress(self, value):
        self.progress_bar.setValue(value)

    def processing_finished(self, output_path):
        self.progress_bar.setVisible(False)
        self.process_btn.setEnabled(True)
        QMessageBox.information(self, "完成", f"Logo添加完成！\n输出文件: {output_path}")

    def processing_error(self, error_message):
        self.progress_bar.setVisible(False)
        self.process_btn.setEnabled(True)
        QMessageBox.critical(self, "错误", error_message) 